package jsr292.weaver;

import java.io.PrintWriter;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

import jsr292.weaver.opt.Optimizer;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.util.CheckClassAdapter;
import org.objectweb.asm.util.TraceClassVisitor;

public class Agent {
    @SuppressWarnings("serial")
    static class StopError extends Error {
        // nothing
    }
    static StopError STOP_ERROR = new StopError();
    
    private static final boolean ENABLE_RETRANSFORM = true;
    static final boolean ENABLE_CLASS_VERSION_CHECK;
    static final boolean ENABLE_WEAVER;
    static {
        ENABLE_WEAVER = Boolean.parseBoolean(System.getProperty("jsr292.backport.weaver", "true"));
        ENABLE_CLASS_VERSION_CHECK = Boolean.parseBoolean(System.getProperty("jsr292.backport.classversioncheck", "true"));
    }
    
    // downgrade version to 1.5
    static class AgentMethodHandleWeaver extends OnlineWeaver {
        public AgentMethodHandleWeaver(ClassVisitor cv, boolean retransform) {
            super(cv, retransform);
        }

        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            // Warning: invokedynamic should be used only in class tagged with version V1_7 
            if (ENABLE_CLASS_VERSION_CHECK && version < Opcodes.V1_7)
                throw STOP_ERROR;
            
            //System.out.println("weave "+name+" "+version);
            
            if (version >= Opcodes.V1_7)
                version = Opcodes.V1_5;
            
            super.visit(version, access, name, signature, superName, interfaces);
        }
    }
    
    public static Instrumentation instrumentation;
    
    // called when installing agent at runtime
    /*public static void agentmain(String agentArgs, Instrumentation instrumentation) {
        Agent.instrumentation = instrumentation;
    }*/
    
    // called by command line
    public static void premain(@SuppressWarnings("unused") String agentArgs, Instrumentation instrumentation) {
        final boolean redefineSupported = ENABLE_RETRANSFORM && instrumentation.isRedefineClassesSupported();
        
        if (!ENABLE_WEAVER && !redefineSupported) {
            System.err.println("agent is useless, no weaving and no optimizer");
            return;
        }
        
        // save instrumentation to be used by the optimizer
        if (redefineSupported) {
            Agent.instrumentation = instrumentation;
        }
      
        ClassFileTransformer weaverTransformer = new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className,
                    Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
                    byte[] classfileBuffer) throws IllegalClassFormatException {
                
                // don't try to transform class in bootstrap classpath
                // commented because JRuby put its classes in bootstrap path
                //if (loader == null)
                //    return null;
                
                if (!ENABLE_WEAVER) {
                    Optimizer.registerClassDefinition(className, classfileBuffer);
                    return null;
                }
                
        	//System.out.println("agent "+className);
                
        	try {
        	    ClassReader reader = new ClassReader(classfileBuffer);
        	    
        	    // Can't use a ClassWriter(reader), because it will copy the constant pool
        	    // with 1.7 specific 
        	    //ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
        	    ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);

        	    try {
        		reader.accept(new AgentMethodHandleWeaver(writer, redefineSupported), ClassReader.SKIP_FRAMES);
        	    } catch(StopError error) {
        		// class with version less than 1.7, no transformation needed
        		return null;
        	    }

        	    byte[] byteArray = writer.toByteArray();
        	    //new ClassReader(byteArray).accept(new TraceClassVisitor(new PrintWriter(System.out)), 0);
        	    //CheckClassAdapter.verify(new ClassReader(byteArray), loader, true, new PrintWriter(System.out));

        	    if (redefineSupported) {
        		Optimizer.registerClassDefinition(className, byteArray);
        	    }
        	    return byteArray;
        	    
        	} catch(Throwable t) {
        	    System.err.println("JSR292 Backport Agent Error:");
        	    t.printStackTrace();
        	    throw (IllegalClassFormatException)new IllegalClassFormatException().initCause(t);
        	}
            }
        };
        
        ClassFileTransformer optimizerTransformer = new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className,
                    Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
                    byte[] classfileBuffer) throws IllegalClassFormatException {
                
                if (classBeingRedefined == null) {
                    // just skip it, this allow to skip generated trampoline class
                    return null;
                }

                //System.out.println("Optimizer retransform "+className);
                try {
                    byte[] byteArray = Optimizer.reweave(classBeingRedefined, classfileBuffer);
                    //if (byteArray != null)
                    //    CheckClassAdapter.verify(new ClassReader(byteArray), loader, true, new PrintWriter(System.out));
                    return byteArray;
                } catch(Throwable t) {
                    t.printStackTrace();
                    throw (IllegalClassFormatException)new IllegalClassFormatException().initCause(t); 
                }
            }
        };
        
        if (!Capabilities.canRetransform()) {
            instrumentation.addTransformer(weaverTransformer);
            if (redefineSupported) {
                instrumentation.addTransformer(optimizerTransformer);
            }
        } else {
            Capabilities.addTransformer(instrumentation, weaverTransformer, false);
            if (redefineSupported) {
                Capabilities.addTransformer(instrumentation, optimizerTransformer, true);
            }

        }
    }
}
